#!/usr/bin/env python3
import cv2
import numpy as np
import time

import rclpy
from rclpy.node import Node
from rclpy.callback_groups import MutuallyExclusiveCallbackGroup
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from rcl_interfaces.msg import ParameterDescriptor, ParameterType

from ball_collector_interfaces.msg import BallInfo
from ball_collector_interfaces.srv import ChangeColor

# 导入MIPI相机模块
from hobot_vio import libsrcampy as srcampy

class BallDetectionNode(Node):
    """
    使用OpenCV颜色过滤进行球体检测的ROS2节点，直接使用MIPI相机
    """
    def __init__(self):
        super().__init__('ball_detection_node')
        
        # 创建不同的回调组确保线程安全
        self.camera_cb_group = MutuallyExclusiveCallbackGroup()
        self.service_cb_group = MutuallyExclusiveCallbackGroup()
        
        # 声明并获取参数
        self.declare_parameters(
            namespace='',
            parameters=[
                ('target_color', 'red',
                 ParameterDescriptor(type=ParameterType.PARAMETER_STRING,
                                     description='Target color to detect')),
                ('detection_freq', 10.0,
                 ParameterDescriptor(type=ParameterType.PARAMETER_DOUBLE,
                                     description='Detection frequency in Hz')),
                ('min_ball_radius', 10,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='Minimum ball radius in pixels')),
                ('max_ball_radius', 300,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='Maximum ball radius in pixels')),
                ('min_circularity', 0.6,
                 ParameterDescriptor(type=ParameterType.PARAMETER_DOUBLE,
                                     description='Minimum circularity (0-1)')),
                ('debug_view', True,
                 ParameterDescriptor(type=ParameterType.PARAMETER_BOOL,
                                     description='Publish debug visualization')),
                ('camera_width', 640,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='Camera width in pixels')),
                ('camera_height', 480,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='Camera height in pixels')),
                ('sensor_width', 1920,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='Sensor width in pixels')),
                ('sensor_height', 1080,
                 ParameterDescriptor(type=ParameterType.PARAMETER_INTEGER,
                                     description='Sensor height in pixels'))
            ]
        )
        
        # 获取参数值
        self.target_color = self.get_parameter('target_color').value
        self.detection_freq = self.get_parameter('detection_freq').value
        self.min_ball_radius = self.get_parameter('min_ball_radius').value
        self.max_ball_radius = self.get_parameter('max_ball_radius').value
        self.min_circularity = self.get_parameter('min_circularity').value
        self.debug_view = self.get_parameter('debug_view').value
        self.camera_width = self.get_parameter('camera_width').value
        self.camera_height = self.get_parameter('camera_height').value
        self.sensor_width = self.get_parameter('sensor_width').value
        self.sensor_height = self.get_parameter('sensor_height').value
        
        # 添加参数变化回调
        self.add_on_set_parameters_callback(self.parameters_callback)
        
        # 设置颜色范围 (HSV)
        self.color_ranges = {
            'red': [
                (np.array([0, 100, 100]), np.array([10, 255, 255])),
                (np.array([160, 100, 100]), np.array([180, 255, 255]))  # 红色在HSV空间中分为两部分
            ],
            'blue': [(np.array([100, 100, 100]), np.array([140, 255, 255]))],
            'black': [(np.array([0, 0, 0]), np.array([180, 100, 50]))],
            'yellow': [(np.array([20, 100, 100]), np.array([30, 255, 255]))],
            'green': [(np.array([40, 100, 100]), np.array([80, 255, 255]))],
            'purple': [(np.array([140, 100, 100]), np.array([160, 255, 255]))],
            'orange': [(np.array([10, 100, 100]), np.array([20, 255, 255]))],
        }
        
        # 初始化CV桥接器
        self.bridge = CvBridge()
        
        # 变量初始化
        self.current_frame = None
        self.frame_id = 'camera_frame'
        self.detection_time = 0.0
        
        # 初始化MIPI相机
        self.init_camera()
        
        # 创建发布者
        # 发布检测到的球信息
        self.ball_pub = self.create_publisher(BallInfo, 'detected_ball', 10)
        
        # 发布处理后的图像（如果启用了调试视图）
        if self.debug_view:
            self.debug_image_pub = self.create_publisher(
                Image,
                'detection_image',
                10
            )
            
            # 添加原始图像发布者，用于调试
            self.raw_image_pub = self.create_publisher(
                Image,
                'raw_image',
                10
            )
        
        # 创建服务
        self.change_color_srv = self.create_service(
            ChangeColor,
            'change_target_color',
            self.change_color_callback,
            callback_group=self.service_cb_group
        )
        
        # 创建定时器
        self.detection_timer = self.create_timer(
            1.0 / self.detection_freq,
            self.detection_callback
        )
        
        # 性能统计
        self.frame_count = 0
        self.total_detection_time = 0.0
        
        self.get_logger().info(f"球体检测节点已初始化，目标颜色: {self.target_color}")
    
    def init_camera(self):
        """初始化MIPI相机"""
        try:
            # 获取适合显示的分辨率
            self.disp_w, self.disp_h = self.get_display_res()
            
            # 创建Camera对象
            self.cam = srcampy.Camera()
            
            # 打开相机，使用与mipi_camera.py相同的参数
            self.cam.open_cam(0, -1, -1, 
                              [self.camera_width, self.disp_w], 
                              [self.camera_height, self.disp_h],
                              self.sensor_height, self.sensor_width)
            
            self.get_logger().info(f"MIPI相机初始化成功，分辨率: {self.camera_width}x{self.camera_height}")
            
            # 创建相机更新定时器，比检测频率更高
            self.camera_timer = self.create_timer(
                0.03,  # 约30FPS
                self.update_camera_frame
            )
            
        except Exception as e:
            self.get_logger().error(f"MIPI相机初始化失败: {e}")
            raise
    
    def get_display_res(self):
        """获取适合显示的分辨率，从mipi_camera.py移植"""
        disp_w_small = 640
        disp_h_small = 480
        
        try:
            disp = srcampy.Display()
            resolution_list = disp.get_display_res()
            
            if (self.sensor_width, self.sensor_height) in resolution_list:
                self.get_logger().info(f"分辨率 {self.sensor_width}x{self.sensor_height} 支持显示")
                return int(self.sensor_width), int(self.sensor_height)
            else:
                self.get_logger().info(f"分辨率 {self.sensor_width}x{self.sensor_height} 不在支持列表中，寻找合适的分辨率")
                for res in resolution_list:
                    # 排除0分辨率
                    if res[0] == 0 and res[1] == 0:
                        break
                    else:
                        disp_w_small = res[0]
                        disp_h_small = res[1]

                    # 找到不大于传感器分辨率的最大分辨率
                    if res[0] <= self.sensor_width and res[1] <= self.sensor_height:
                        self.get_logger().info(f"选择分辨率 {res[0]}x{res[1]}")
                        return int(res[0]), int(res[1])
            
            disp.close()
        except Exception as e:
            self.get_logger().error(f"获取显示分辨率失败: {e}")
        
        return disp_w_small, disp_h_small
    
    def update_camera_frame(self):
        """从MIPI相机获取图像帧"""
        try:
            # 获取相机图像
            img_data = self.cam.get_img(2, self.camera_width, self.camera_height)
            
            # 转换为numpy数组
            img_np = np.frombuffer(img_data, dtype=np.uint8)
            
            # 重塑为图像形状 (NV12格式)
            # NV12格式: 先是Y平面，后是交错的UV平面
            y_size = self.camera_width * self.camera_height
            y_plane = img_np[:y_size].reshape(self.camera_height, self.camera_width)
            
            # 创建与Y平面相同大小的U和V平面
            u_plane = np.zeros((self.camera_height, self.camera_width), dtype=np.uint8)
            v_plane = np.zeros((self.camera_height, self.camera_width), dtype=np.uint8)
            
            # 从交错的UV平面中提取U和V
            if len(img_np) > y_size:  # 确保有足够的数据包含UV平面
                uv_plane = img_np[y_size:].reshape(self.camera_height//2, self.camera_width//2, 2)
                
                # 放大UV平面到与Y平面相同的大小（简单的重复）
                u_plane[::2, ::2] = uv_plane[:, :, 0]
                u_plane[1::2, ::2] = uv_plane[:, :, 0]
                u_plane[::2, 1::2] = uv_plane[:, :, 0]
                u_plane[1::2, 1::2] = uv_plane[:, :, 0]
                
                v_plane[::2, ::2] = uv_plane[:, :, 1]
                v_plane[1::2, ::2] = uv_plane[:, :, 1]
                v_plane[::2, 1::2] = uv_plane[:, :, 1]
                v_plane[1::2, 1::2] = uv_plane[:, :, 1]
            
            # 合并YUV通道并转换为BGR格式
            yuv = cv2.merge([y_plane, u_plane, v_plane])
            bgr = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
            
            # 更新当前帧
            self.current_frame = bgr
            
            # 发布原始图像用于调试
            if self.debug_view:
                try:
                    img_msg = self.bridge.cv2_to_imgmsg(bgr, "bgr8")
                    img_msg.header.stamp = self.get_clock().now().to_msg()
                    img_msg.header.frame_id = self.frame_id
                    self.raw_image_pub.publish(img_msg)
                except Exception as e:
                    self.get_logger().error(f"发布原始图像失败: {e}")
            
        except Exception as e:
            self.get_logger().error(f"获取相机图像失败: {e}")
    
    def detection_callback(self):
        """执行周期性检测"""
        if self.current_frame is None:
            self.get_logger().warn("没有可用的图像帧")
            return
        
        start_time = time.time()
        frame = self.current_frame.copy()
        
        # 存储检测结果
        balls = []
        
        # 使用颜色过滤检测目标颜色的球
        color_balls = self.detect_balls_by_color(frame, self.target_color)
        
        # 将结果添加到球列表
        balls.extend(color_balls)
        
        # 发布检测结果
        if balls:
            # 如果发现多个球，选择最大的一个（通常是最近的）
            if len(balls) > 1:
                # 按直径大小排序
                balls.sort(key=lambda x: x['diameter'], reverse=True)
            
            # 发布球信息
            self.publish_ball_info(balls[0], frame)
        elif self.debug_view and frame is not None:
            # 如果未检测到球但启用了调试视图，仍然发布处理后的图像
            try:
                # 添加提示文本
                cv2.putText(frame, "No balls detected", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                
                # 添加目标颜色信息
                cv2.putText(frame, f"Target: {self.target_color}", (10, 60),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                
                # 发布调试图像
                img_msg = self.bridge.cv2_to_imgmsg(frame, "bgr8")
                img_msg.header.stamp = self.get_clock().now().to_msg()
                img_msg.header.frame_id = self.frame_id
                self.debug_image_pub.publish(img_msg)
            except Exception as e:
                self.get_logger().error(f"发布调试图像失败: {e}")
        
        # 计算检测时间
        self.detection_time = time.time() - start_time
        
        # 记录性能
        self.frame_count += 1
        self.total_detection_time += self.detection_time
        
        if self.frame_count % 30 == 0:  # 每30帧记录一次
            avg_time = self.total_detection_time / 30
            fps = 30 / self.total_detection_time if self.total_detection_time > 0 else 0.0
            self.get_logger().debug(f"平均检测时间: {avg_time:.3f}秒, FPS: {fps:.2f}")
            self.total_detection_time = 0.0
    
    def detect_balls_by_color(self, frame, target_color):
        """使用颜色过滤检测指定颜色的球"""
        balls = []
        
        if target_color not in self.color_ranges:
            self.get_logger().warn(f"未知颜色: {target_color}")
            return balls
        
        try:
            # 转换到HSV颜色空间
            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
            
            # 创建掩码
            mask = np.zeros(frame.shape[:2], dtype=np.uint8)
            
            # 应用颜色范围
            for lower, upper in self.color_ranges[target_color]:
                color_mask = cv2.inRange(hsv, lower, upper)
                mask = cv2.bitwise_or(mask, color_mask)
            
            # 应用形态学操作以减少噪声
            kernel = np.ones((5, 5), np.uint8)
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)  # 开运算（先腐蚀后膨胀）
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel) # 闭运算（先膨胀后腐蚀）
            
            # 查找轮廓
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            # 在调试视图上显示掩码
            if self.debug_view:
                mask_color = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
                # 在帧的角落展示掩码，缩小到原尺寸的1/4
                h, w = mask.shape
                small_mask = cv2.resize(mask_color, (w//4, h//4))
                frame[0:h//4, 0:w//4] = small_mask
            
            # 处理每个轮廓
            for contour in contours:
                # 计算面积
                area = cv2.contourArea(contour)
                
                # 过滤小区域
                if area < np.pi * (self.min_ball_radius ** 2):
                    continue
                
                # 拟合圆
                (x, y), radius = cv2.minEnclosingCircle(contour)
                center = (int(x), int(y))
                radius = int(radius)
                
                # 检查半径范围
                if radius < self.min_ball_radius or radius > self.max_ball_radius:
                    continue
                
                # 计算圆度 (4*pi*面积/周长^2)
                perimeter = cv2.arcLength(contour, True)
                circularity = 4 * np.pi * area / (perimeter ** 2) if perimeter > 0 else 0
                
                # 过滤非圆形对象
                if circularity < self.min_circularity:
                    continue
                
                # 创建球信息
                ball_info = {
                    'center': center,
                    'diameter': 2 * radius,
                    'confidence': circularity,  # 使用圆度作为置信度
                    'color': target_color,
                    'is_target': True,  # 因为我们只检测目标颜色
                    'bbox': (int(x - radius), int(y - radius), int(x + radius), int(y + radius))
                }
                
                balls.append(ball_info)
        
        except Exception as e:
            self.get_logger().error(f"颜色检测异常: {e}")
        
        return balls
    
    def publish_ball_info(self, ball, frame):
        """发布球信息"""
        # 创建并发布球信息消息
        ball_msg = BallInfo()
        ball_msg.header.stamp = self.get_clock().now().to_msg()
        ball_msg.header.frame_id = self.frame_id
        ball_msg.id = 0  # 单个球，ID为0
        ball_msg.color = ball['color']
        ball_msg.center.x = float(ball['center'][0])
        ball_msg.center.y = float(ball['center'][1])
        ball_msg.center.z = 0.0
        ball_msg.diameter = float(ball['diameter'])
        ball_msg.confidence = ball['confidence']
        ball_msg.is_target = ball['is_target']
        
        self.ball_pub.publish(ball_msg)
        
        # 在调试图像上标记
        if self.debug_view:
            x1, y1, x2, y2 = ball['bbox']
            
            # 绘制边界框
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            
            # 绘制圆
            cv2.circle(frame, ball['center'], int(ball['diameter'] / 2), (0, 255, 0), 2)
            
            # 添加颜色和置信度标签
            label = f"{ball['color']}: {ball['confidence']:.2f}"
            cv2.putText(frame, label, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            # 绘制中心点
            cv2.circle(frame, ball['center'], 5, (0, 0, 255), -1)
            
            # 添加性能信息
            fps_text = f"FPS: {1.0/self.detection_time:.1f}" if self.detection_time > 0 else "FPS: -"
            cv2.putText(frame, fps_text, (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            
            # 添加目标颜色信息
            cv2.putText(frame, f"Target: {self.target_color}", (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
            
            # 转换回ROS图像消息并发布
            img_msg = self.bridge.cv2_to_imgmsg(frame, "bgr8")
            img_msg.header.stamp = self.get_clock().now().to_msg()
            img_msg.header.frame_id = self.frame_id
            self.debug_image_pub.publish(img_msg)
    
    def change_color_callback(self, request, response):
        """更改目标颜色的服务回调"""
        if request.new_color in self.color_ranges:
            old_color = self.target_color
            self.target_color = request.new_color
            self.get_logger().info(f"目标颜色已更改: {old_color} -> {self.target_color}")
            response.success = True
        else:
            self.get_logger().warn(f"不支持的颜色: {request.new_color}")
            response.success = False
        
        response.current_color = self.target_color
        return response
    
    def parameters_callback(self, params):
        """参数更新回调"""
        for param in params:
            if param.name == 'target_color' and param.type_ == param.Type.STRING:
                if param.value in self.color_ranges:
                    self.target_color = param.value
                    self.get_logger().info(f"目标颜色参数已更新: {self.target_color}")
                else:
                    self.get_logger().warn(f"不支持的颜色参数: {param.value}")
                    return False
            
            elif param.name == 'min_ball_radius' and param.type_ == param.Type.INTEGER:
                if param.value > 0:
                    self.min_ball_radius = param.value
                    self.get_logger().info(f"最小球半径参数已更新: {self.min_ball_radius}")
                else:
                    self.get_logger().warn(f"无效的最小球半径: {param.value}")
                    return False
            
            elif param.name == 'max_ball_radius' and param.type_ == param.Type.INTEGER:
                if param.value > self.min_ball_radius:
                    self.max_ball_radius = param.value
                    self.get_logger().info(f"最大球半径参数已更新: {self.max_ball_radius}")
                else:
                    self.get_logger().warn(f"无效的最大球半径: {param.value}")
                    return False
            
            elif param.name == 'min_circularity' and param.type_ == param.Type.DOUBLE:
                if 0.0 < param.value <= 1.0:
                    self.min_circularity = param.value
                    self.get_logger().info(f"最小圆度参数已更新: {self.min_circularity}")
                else:
                    self.get_logger().warn(f"无效的最小圆度: {param.value}")
                    return False
        
        return True
    
    def destroy_node(self):
        """清理资源"""
        # 关闭相机
        if hasattr(self, 'cam'):
            try:
                self.cam.close_cam()
                self.get_logger().info("已关闭MIPI相机")
            except Exception as e:
                self.get_logger().error(f"关闭相机失败: {e}")
        
        super().destroy_node()

def main(args=None):
    rclpy.init(args=args)
    
    try:
        # 创建节点
        node = BallDetectionNode()
        
        try:
            rclpy.spin(node)
        except KeyboardInterrupt:
            pass
        finally:
            node.destroy_node()
    except Exception as e:
        print(f"初始化节点失败: {e}")
    finally:
        rclpy.shutdown()

if __name__ == '__main__':
    main()